import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import numpy as np
from .Loss import Loss


class SUKE_LOSS(Loss):

    def __init__(self, adv_temperature=None, margin=6.0):
        super(SUKE_LOSS, self).__init__()
        self.margin = nn.Parameter(torch.Tensor([margin]))
        self.margin.requires_grad = False
        if adv_temperature != None:
            self.adv_temperature = nn.Parameter(torch.Tensor([adv_temperature]))
            self.adv_temperature.requires_grad = False
            self.adv_flag = True
        else:
            self.adv_flag = False

    def get_weights(self, n_score):
        return F.softmax(-n_score * self.adv_temperature, dim=-1).detach()

    def forward(self, p_score, n_score):
        q_stru_p = p_score[0]
        q_unce_p = p_score[1]
        q_stru_n = n_score[0]
        q_unce_n = n_score[1]
        p_loss = torch.add(q_stru_p, -1).pow(2) + torch.add(q_unce_p, -0.5).pow(2)
        n_loss = q_stru_n.pow(2) + q_unce_n.pow(2)
        loss = p_loss.sum() + n_loss.sum() + self.margin
        return loss

    def predict(self, p_score, n_score):
        score = self.forward(p_score, n_score)
        return score.cpu().data.numpy()